#include <glad/glad.h>
#include <GLFW/glfw3.h>              
#include <iostream>
#include <string.h>
#include "acceleratetransfer.h"
#include "shader.h"

#define GET_STR(x) #x

static bool opengl_initalized = false;

const char* vertxStr = GET_STR(
\#version 330 core\n
layout (location = 0) in vec3 aPos;
layout (location = 1) in vec2 aTexCoord;
out vec2 TexCoord;
uniform mat4 texture_scale;
void main()
{
    gl_Position = vec4(aPos, 1.0);
    vec4 s = texture_scale * vec4(aTexCoord,1.0,1.0);
    TexCoord = vec2(s.x, s.y);
});

const char* fragtxStr = GET_STR(
\#version 330 core\n
out vec4 FragColor;
in vec2 TexCoord;

// texture samplers
uniform sampler2D textureY;
uniform sampler2D textureU;
uniform sampler2D textureV;
void main()
{
    vec3 yuv;
    vec3 rgb;
    yuv.x = texture(textureY,vec2(TexCoord.x, TexCoord.y)).r;
    yuv.y = texture(textureU, vec2(TexCoord.x, TexCoord.y)).r-0.5;
    yuv.z = texture(textureV, vec2(TexCoord.x, TexCoord.y)).r-0.5;
    rgb = mat3(1.0,1.0,1.0,0.0,-0.39465,2.03211,1.13983,-0.58060,0.0)*yuv;
    FragColor = vec4(rgb, 1.0);
});


AccelerateTransfer::AccelerateTransfer()
{

}

AccelerateTransfer::~AccelerateTransfer(){
    glDeleteVertexArrays(1, &VAO);
    glDeleteFramebuffers(1,&FBO);
    glDeleteBuffers(1, &VBO);
    glDeleteBuffers(1, &EBO);
    glDeleteBuffers(6, PBOS);
    glDeleteTextures(1,&texColorBuffer);
    glDeleteTextures(1,&tf[Y].textureId);
    glDeleteTextures(1,&tf[U].textureId);
    glDeleteTextures(1,&tf[V].textureId);
    glDeleteBuffers(2,rgbpbos);
    delete shd;
    if(ydata)
        delete[] ydata;
    if(udata)
        delete[] udata;
    if(vdata)
        delete[] vdata;
    glfwTerminate();
}

bool AccelerateTransfer::transfer_init(int winW,int winH,const char* title){
    win_width = winW;
    win_height = winH;
    window = glfwCreateWindow(win_width, win_height, title, NULL, NULL);
    if (window == NULL)
    {
        std::cout << "Failed to create GLFW window" << std::endl;
        glfwTerminate();
        return false;
    }
    glfwMakeContextCurrent(window);
    int ret = gladLoadGLLoader((GLADloadproc)glfwGetProcAddress);
    if (!ret)
    {
        std::cout << "Failed to initialize GLAD" << std::endl;
        return false;
    }
    shd = new Shader;
    shd->setShaderCode(vertxStr,fragtxStr);
    if(!shd->create_init_shader_program())
        return false;
    else
        return true;
}

void AccelerateTransfer::prepare_context(int w,int h,int pixelBytes,float scalew,float scaleh){
    float vertices[] = {
        // positions        // texture coords
         1.0f,  1.0f, 0.0f, 1.0f, 1.0f, // top right
         1.0f, -1.0f, 0.0f, 1.0f, 0.0f, // bottom right
        -1.0f, -1.0f, 0.0f, 0.0f, 0.0f , // bottom left
        -1.0f,  1.0f, 0.0f, 0.0f, 1.0f   // top left
    };
    unsigned int indices[] = {
        0, 1, 3, // first triangle
        1, 2, 3  // second triangle
    };
    glGenVertexArrays(1, &VAO);
    glGenBuffers(1, &VBO);
    glGenBuffers(1, &EBO);
    glBindVertexArray(VAO);
    glBindBuffer(GL_ARRAY_BUFFER, VBO);
    glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW);

    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, EBO);
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, sizeof(indices), indices, GL_STATIC_DRAW);

    // position attribute
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 5 * sizeof(float), (void*)0);
    glEnableVertexAttribArray(0);
    // texture coord attribute
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 5 * sizeof(float), (void*)(3 * sizeof(float)));
    glEnableVertexAttribArray(1);

    glGenTextures(1, &tf[Y].textureId);
    glBindTexture(GL_TEXTURE_2D, tf[Y].textureId);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_BORDER);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_BORDER);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RED, w, h, 0, GL_RED, GL_UNSIGNED_BYTE, NULL);
    glGenerateMipmap(GL_TEXTURE_2D);
    tf[Y].pixelFormatCode = GL_RED;
    tf[Y].pixelTypeCode = GL_UNSIGNED_BYTE;
    tf[Y].textureH = h;
    tf[Y].textureW = w;

    glGenTextures(1, &tf[U].textureId);
    glBindTexture(GL_TEXTURE_2D, tf[U].textureId);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_BORDER);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_BORDER);

    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RED, w/2, h/2, 0, GL_RED, GL_UNSIGNED_BYTE, NULL);
    glGenerateMipmap(GL_TEXTURE_2D);
    tf[U].pixelFormatCode = GL_RED;
    tf[U].pixelTypeCode = GL_UNSIGNED_BYTE;
    tf[U].textureH = h/2;
    tf[U].textureW = w/2;

    glGenTextures(1, &tf[V].textureId);
    glBindTexture(GL_TEXTURE_2D, tf[V].textureId);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_BORDER);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_BORDER);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RED, w/2, h/2, 0, GL_RED, GL_UNSIGNED_BYTE, NULL);
    glGenerateMipmap(GL_TEXTURE_2D);
    tf[V].pixelFormatCode = GL_RED;
    tf[V].pixelTypeCode = GL_UNSIGNED_BYTE;
    tf[V].textureH = h/2;
    tf[V].textureW = w/2;

    shd->use();
    glUniform1i(glGetUniformLocation(shd->ID, "textureY"), Y);
    glUniform1i(glGetUniformLocation(shd->ID, "textureU"), U);
    glUniform1i(glGetUniformLocation(shd->ID, "textureV"), V);
    
    glGenFramebuffers(1, &FBO);
    glBindFramebuffer(GL_FRAMEBUFFER, FBO);//绑定帧缓冲区到自定义帧缓冲对象FBO
    glGenTextures(1, &texColorBuffer);
    glBindTexture(GL_TEXTURE_2D, texColorBuffer);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, win_width, win_height, 0, GL_RGB, GL_UNSIGNED_BYTE, NULL);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);

    glBindTexture(GL_TEXTURE_2D, 0);
    // 将texColorBuffer附加到当前绑定着的帧缓冲对象FBO
    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, texColorBuffer, 0);
    if (glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE)
        std::cout << "ERROR::FRAMEBUFFER:: Framebuffer is not complete!" << std::endl;
    glBindFramebuffer(GL_FRAMEBUFFER, FBO);//绑定帧缓冲区到屏幕缓冲对象（屏幕缓冲对象为默认id 0）

    if(pixelBytes==4){
        readPixelType = GL_FLOAT;
        readPixelSize = sizeof(GL_FLOAT);
    }

    float trans_scale[16]={scalew,0.0,0.0,0.0,
                           0.0,scaleh,0.0,0.0,
                           0.0,0.0,1.0,0.0,
                           0.0,0.0,0.0,1.0};
    glUniformMatrix4fv(glGetUniformLocation(shd->ID, "texture_scale"), 1, GL_FALSE, trans_scale);

    glGenBuffers(6, PBOS);
    glBindBuffer(GL_PIXEL_PACK_BUFFER, PBOS[0]); //绑定 R0
    glBufferData(GL_PIXEL_PACK_BUFFER, win_width * win_height*readPixelSize, 0, GL_STREAM_READ);
    glBindBuffer(GL_PIXEL_PACK_BUFFER, PBOS[1]); //绑定 G0
    glBufferData(GL_PIXEL_PACK_BUFFER, win_width* win_height*readPixelSize, 0, GL_STREAM_READ);
    glBindBuffer(GL_PIXEL_PACK_BUFFER, PBOS[2]); //绑定 B0
    glBufferData(GL_PIXEL_PACK_BUFFER, win_width* win_height*readPixelSize, 0, GL_STREAM_READ);

    glBindBuffer(GL_PIXEL_PACK_BUFFER, PBOS[3]); //绑定 R1
    glBufferData(GL_PIXEL_PACK_BUFFER, win_width * win_height *readPixelSize, 0, GL_STREAM_READ);
    glBindBuffer(GL_PIXEL_PACK_BUFFER, PBOS[4]); //绑定 G1
    glBufferData(GL_PIXEL_PACK_BUFFER, win_width * win_height *readPixelSize, 0, GL_STREAM_READ);
    glBindBuffer(GL_PIXEL_PACK_BUFFER, PBOS[5]); //绑定 B1
    glBufferData(GL_PIXEL_PACK_BUFFER, win_width * win_height *readPixelSize, 0, GL_STREAM_READ);

    glGenBuffers(2, rgbpbos);
    glBindBuffer(GL_PIXEL_PACK_BUFFER, rgbpbos[0]); //绑定 rgb0
    glBufferData(GL_PIXEL_PACK_BUFFER, win_width * win_height *3, 0, GL_STREAM_READ);
    glBindBuffer(GL_PIXEL_PACK_BUFFER, rgbpbos[1]); //绑定 rgb1
    glBufferData(GL_PIXEL_PACK_BUFFER, win_width * win_height *3, 0, GL_STREAM_READ);

    glBindBuffer(GL_PIXEL_PACK_BUFFER, 0); //解绑
    glEnable(GL_UNSIGNED_NORMALIZED);
    glEnable(GL_SIGNED_NORMALIZED);
    ydata = new unsigned char[w*h];
    udata = new unsigned char[w*h/4];
    vdata = new unsigned char[w*h/4];
    memset(ydata,0,w*h);
    memset(udata,0,w*h/4);
    memset(vdata,0,w*h/4);
}

void AccelerateTransfer::update_texture(YUVCHANNEL channel_yuv, const unsigned char* data){
    if(channel_yuv == Y && data)
        memcpy(ydata, data, tf[Y].textureH*tf[Y].textureW);
    if(channel_yuv == U && data)
        memcpy(udata, data, tf[U].textureH*tf[U].textureW);
    if(channel_yuv == V && data)
        memcpy(vdata, data, tf[V].textureH*tf[V].textureW);
}

void AccelerateTransfer::draw_once(){
    glClearColor(1.0f, 1.0f, 1.0f, 1.0f);
    glClear(GL_COLOR_BUFFER_BIT);
    shd->use();
    glPixelStorei(GL_UNPACK_ALIGNMENT, 1);
    glActiveTexture(GL_TEXTURE0);
    glBindTexture(GL_TEXTURE_2D, tf[Y].textureId);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RED, tf[Y].textureW, tf[Y].textureH, 0, GL_RED, GL_UNSIGNED_BYTE, ydata);
    glActiveTexture(GL_TEXTURE1);
    glBindTexture(GL_TEXTURE_2D, tf[U].textureId);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RED, tf[U].textureW, tf[U].textureH, 0, GL_RED, GL_UNSIGNED_BYTE, udata);
    glActiveTexture(GL_TEXTURE2);
    glBindTexture(GL_TEXTURE_2D, tf[V].textureId);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RED, tf[V].textureW, tf[V].textureH, 0, GL_RED, GL_UNSIGNED_BYTE, vdata);
    glPixelStorei(GL_UNPACK_ALIGNMENT, 4);
    //glm::mat4 trans = glm::mat4(1.0);
    //trans = glm::scale(trans, glm::vec3(1.17 , 1.0, 1.0));
    //glUniformMatrix4fv(glGetUniformLocation(shd->ID, "texture_scale"), 1, GL_FALSE, glm::value_ptr(trans));
    glBindVertexArray(VAO);
    glDrawElements(GL_TRIANGLES, 6, GL_UNSIGNED_INT, 0);

    //读取帧缓冲区中的渲染结果到GL_PIXEL_PACK_BUFFER
    glGetIntegerv(GL_RENDERBUFFER_WIDTH,0);
    glBindBuffer(GL_PIXEL_PACK_BUFFER, PBOS[pboGroupIdxCur * 3 + R]);
    glReadPixels(0,0, win_width, win_height, GL_RED, readPixelType, 0);
    glBindBuffer(GL_PIXEL_PACK_BUFFER, PBOS[pboGroupIdxCur * 3 + G]);
    glReadPixels(0,0, win_width, win_height, GL_GREEN, readPixelType, 0);
    glBindBuffer(GL_PIXEL_PACK_BUFFER, PBOS[pboGroupIdxCur * 3 + B]);
    glReadPixels(0,0, win_width, win_height, GL_BLUE, readPixelType, 0);

    glBindBuffer(GL_PIXEL_PACK_BUFFER, rgbpbos[pboGroupIdxCur]);
    glReadPixels(0,0, win_width, win_height, GL_RGB, GL_UNSIGNED_BYTE, 0);
    pboGroupIdxCur ^= pboGroupIdxNext;
    pboGroupIdxNext ^= pboGroupIdxCur;
    pboGroupIdxCur ^= pboGroupIdxNext;
}

int AccelerateTransfer::read_pixel_channel(RGBCHANNEL channel_rgb, const int planesize, int& w,int& h, unsigned char* plane){
    if(channel_rgb<0||channel_rgb>2)
        return 0;
    int src_h = win_height;
    int src_w = win_width;
    if(plane==nullptr||planesize<readPixelSize*src_h*src_w)
        return 0;
    glBindBuffer(GL_PIXEL_PACK_BUFFER, PBOS[pboGroupIdxNext * 3 + channel_rgb]);
    void* pic = glMapBuffer(GL_PIXEL_PACK_BUFFER, GL_READ_ONLY);
    memcpy(plane, pic, src_w * src_h*readPixelSize);
    glUnmapBuffer(GL_PIXEL_PACK_BUFFER);
    w = src_w;
    h = src_h;
    return w*h*readPixelSize;
}

int AccelerateTransfer::read_pixel_rgb(const int planesize, int& w,int& h, unsigned char* plane){
    int src_h = win_height;
    int src_w = win_width;
    if(plane==nullptr||planesize<3*src_h*src_w){
        return 0;
    }
    glBindBuffer(GL_PIXEL_PACK_BUFFER, rgbpbos[pboGroupIdxNext]);
    void* pic = glMapBuffer(GL_PIXEL_PACK_BUFFER, GL_READ_ONLY);
    memcpy(plane, pic, src_w * src_h*3);
    glUnmapBuffer(GL_PIXEL_PACK_BUFFER);
    w = src_w;
    h = src_h;
    return w*h*3;
}


bool AccelerateTransfer::opengl_init(int major_version, int minor_version,bool windowVisable){
    if(opengl_initalized)return opengl_initalized;
    if (!glfwInit()) return false;
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, major_version);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, minor_version);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
    if(!windowVisable)
        glfwWindowHint(GLFW_VISIBLE, GL_FALSE);//不显示窗口(比如离屏渲染)
    opengl_initalized = true;
    return opengl_initalized;
}
